#include "VulkanApplication.h"
#include <iostream>

VulkanApplication::VulkanApplication(const char* appName, int w, int h)
    : _instance(nullptr), _physicalDevice(nullptr), _logicDevice(nullptr),
      _surface(nullptr), _swapChain(nullptr), _commandPool(nullptr),
      _commandBuffer(nullptr), _renderPass(nullptr), _window(nullptr)
{
    if (!glfwInit())
    {
        printf("Failed to initialize GLFW.\n");
        _initialized = false; return;
    }

    if (!glfwVulkanSupported())
    {
        printf("Vulkan not supported.\n");
        _initialized = false; return;
    }

    _initialized = createInstance(appName);
    if (_initialized) _initialized = obtainPhysicalDevice();
    if (_initialized) _initialized = createLogicDevice();
    if (_initialized) _initialized = createWindowSurface(appName, w, h);
    if (_initialized) _initialized = createSwapchain();
    if (_initialized) _initialized = createImageViews();
    if (_initialized) _initialized = createCommandBuffer();
    if (_initialized) _initialized = createRenderPass();
    if (_initialized) _initialized = createFrameBuffers(w, h);
}

VulkanApplication::~VulkanApplication()
{
    cleanupSwapchain();
    for (std::map<std::string, VkSemaphore>::iterator it = _semaphores.begin();
         it != _semaphores.end(); ++it) vkDestroySemaphore(_logicDevice, it->second, nullptr);
    for (std::map<std::string, VkFence>::iterator it = _fences.begin();
         it != _fences.end(); ++it) vkDestroyFence(_logicDevice, it->second, nullptr);

    if (_commandPool != nullptr) vkDestroyCommandPool(_logicDevice, _commandPool, nullptr);
    if (_renderPass != nullptr) vkDestroyRenderPass(_logicDevice, _renderPass, nullptr);
    if (_surface != nullptr) vkDestroySurfaceKHR(_instance, _surface, nullptr);
    if (_window != nullptr) glfwDestroyWindow(_window);
    if (_logicDevice != nullptr) vkDestroyDevice(_logicDevice, nullptr);
    if (_instance != nullptr) vkDestroyInstance(_instance, nullptr);
    glfwTerminate();
}

void VulkanApplication::mouseButtonCallback(GLFWwindow* window, int button, int action, int mods)
{
    printf("MOUSEBUTTON: %d, %d\n", button, action);
}

void VulkanApplication::keyCallback(GLFWwindow* window, int key, int scancode, int action, int mods)
{
    printf("KEY: %d, %d\n", key, action);
    if (key == GLFW_KEY_ESCAPE) glfwSetWindowShouldClose(window, 1);
}

void VulkanApplication::resizeCallback(GLFWwindow* window, int width, int height)
{
    printf("Resized: %d x %d\n", width, height);
}

bool VulkanApplication::createInstance(const char* appName)
{
    uint32_t glfwExtCount = 0;
    const char** glfwExtensions = glfwGetRequiredInstanceExtensions(&glfwExtCount);
    printf("GLFW initialized. It requires the following extensions:\n");
    for (uint32_t i = 0; i < glfwExtCount; i++) printf("    %s\n", glfwExtensions[i]);
    
    const char* layerNames[] =
    {
        "VK_LAYER_KHRONOS_validation"
    };

    VkApplicationInfo appInfo
    {
        VK_STRUCTURE_TYPE_APPLICATION_INFO,
        nullptr,
        appName,
        VK_MAKE_VERSION(1, 0, 0),
        "VulkanClass",
        VK_MAKE_VERSION(1, 0, 0),
        VK_API_VERSION_1_3,
    };

    VkInstanceCreateInfo instanceInfo
    {
        VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO,
        nullptr, 0,
        &appInfo,
        1, layerNames,
        glfwExtCount, glfwExtensions
    };

    VkResult result = vkCreateInstance(&instanceInfo, nullptr, &_instance);
    if (result != VK_SUCCESS)
    {
        printf("Vulkan instance creation failed: %d\n", result);
        return false;
    }
    return true;
}

bool VulkanApplication::obtainPhysicalDevice()
{
    uint32_t deviceCount = 0;
    VkResult result = vkEnumeratePhysicalDevices(_instance, &deviceCount, nullptr);
    if (result != VK_SUCCESS || deviceCount == 0)
    {
        printf("Failed to enumerate physical devices.\n");
        return false;
    }

    std::vector<VkPhysicalDevice> devices(deviceCount);
    vkEnumeratePhysicalDevices(_instance, &deviceCount, devices.data());
    printf("vkEnumeratePhysicalDevices() has found some physical devices:\n");
    for (uint32_t i = 0; i < deviceCount; ++i)
    {
        VkPhysicalDeviceProperties props;
        vkGetPhysicalDeviceProperties(devices[i], &props);

        if (props.deviceType == VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU)
            printf("    %d: Discrete GPU, %s\n", i, props.deviceName);
        else if (props.deviceType == VK_PHYSICAL_DEVICE_TYPE_INTEGRATED_GPU)
            printf("    %d: Integrated GPU, %s\n", i, props.deviceName);
        else
            printf("    %d: Other GPU, %s\n", i, props.deviceName);
    }

    _physicalDevice = devices[0];
    return true;
}

bool VulkanApplication::createLogicDevice()
{
    uint32_t queueFamilyCount = 0;
    vkGetPhysicalDeviceQueueFamilyProperties(_physicalDevice, &queueFamilyCount, nullptr);
    if (queueFamilyCount == 0)
    {
        printf("Failed to get queue family data.\n");
        return false;
    }

    std::vector<VkQueueFamilyProperties> queueFamilies(queueFamilyCount);
    vkGetPhysicalDeviceQueueFamilyProperties(_physicalDevice, &queueFamilyCount, queueFamilies.data());

    VkPhysicalDeviceFeatures features;
    vkGetPhysicalDeviceFeatures(_physicalDevice, &features);

    printf("Current physical device supports:\n");
    printf("    Graphics = %d\n", (queueFamilies[0].queueFlags & VK_QUEUE_GRAPHICS_BIT) > 0);
    printf("    Compute = %d\n", (queueFamilies[0].queueFlags & VK_QUEUE_COMPUTE_BIT) > 0);
    printf("    Geometry = %d\n", features.geometryShader);
    printf("    Tessellation = %d\n", features.tessellationShader);

    char* deviceExtensions[] =
    {
        "VK_KHR_swapchain"
    };
    
    float queuePriorities = 1.0f;
    VkDeviceQueueCreateInfo queueCreateInfo
    {
        VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO,
        nullptr, 0,
        0, 1,  // queueFamilyIndex, queueCount
        &queuePriorities
    };

    VkDeviceCreateInfo createInfo
    {
        VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO,
        nullptr, 0,
        1, &queueCreateInfo,
        0, nullptr,
        1, deviceExtensions,
        &features
    };

    VkResult result = vkCreateDevice(_physicalDevice, &createInfo, nullptr, &_logicDevice);
    if (result != VK_SUCCESS)
    {
        printf("Failed to create logic device.\n");
        return false;
    }
    return true;
}

bool VulkanApplication::createWindowSurface(const char* appName, int w, int h)
{
    glfwWindowHint(GLFW_CLIENT_API, GLFW_NO_API);
    _window = glfwCreateWindow(w, h, appName, nullptr, nullptr);
    if (!_window)
    {
        printf("Failed to create GLFW window.\n");
        return false;
    }

    VkResult result = glfwCreateWindowSurface(_instance, _window, nullptr, &_surface);
    if (result != VK_SUCCESS)
    {
        printf("Failed to create window surface.\n");
        return false;
    }

    glfwSetWindowUserPointer(_window, this);
    glfwSetMouseButtonCallback(_window, [](GLFWwindow* w, int button, int action, int mods)
    {
        VulkanApplication* app = static_cast<VulkanApplication*>(glfwGetWindowUserPointer(w));
        app->mouseButtonCallback(w, button, action, mods);
    });
    glfwSetKeyCallback(_window, [](GLFWwindow* w, int key, int scancode, int action, int mods)
    {
        VulkanApplication* app = static_cast<VulkanApplication*>(glfwGetWindowUserPointer(w));
        app->keyCallback(w, key, scancode, action, mods);
    });
    glfwSetWindowSizeCallback(_window, [](GLFWwindow* w, int width, int height)
    {
        VulkanApplication* app = static_cast<VulkanApplication*>(glfwGetWindowUserPointer(w));
        app->resizeCallback(w, width, height);
    });
    return true;
}

bool VulkanApplication::createSwapchain()
{
    // Get surface capabilities
    VkSurfaceCapabilitiesKHR caps;
    VkResult result = vkGetPhysicalDeviceSurfaceCapabilitiesKHR(_physicalDevice, _surface, &caps);
    if (result != VK_SUCCESS)
    {
        printf("Failed to get surface capabilities.\n");
        return false;
    }

    // Get surface formats and present modes
    uint32_t formatCount = 0, presentModeCount = 0;
    result = vkGetPhysicalDeviceSurfaceFormatsKHR(_physicalDevice, _surface, &formatCount, nullptr);
    if (result != VK_SUCCESS || formatCount == 0)
    {
        printf("Failed to get surface formats.\n");
        return false;
    }

    result = vkGetPhysicalDeviceSurfacePresentModesKHR(
        _physicalDevice, _surface, &presentModeCount, nullptr);
    if (result != VK_SUCCESS || presentModeCount == 0)
    {
        printf("Failed to get surface present modes.\n");
        return false;
    }

    uint32_t formatIndex = 0, presentModeIndex = 0;
    std::vector<VkSurfaceFormatKHR> formats(formatCount);
    std::vector<VkPresentModeKHR> presentModes(presentModeCount);

    vkGetPhysicalDeviceSurfaceFormatsKHR(
        _physicalDevice, _surface, &formatCount, formats.data());
    for (uint32_t i = 0; i < formatCount; ++i)
    {
        //printf("Surface format / colorspace: %d / %d\n", formats[i].format, formats[i].colorSpace);
        if (formats[i].format == VK_FORMAT_B8G8R8A8_SRGB &&
            formats[i].colorSpace == VK_COLOR_SPACE_SRGB_NONLINEAR_KHR) formatIndex = i;
    }

    vkGetPhysicalDeviceSurfacePresentModesKHR(
        _physicalDevice, _surface, &presentModeCount, presentModes.data());
    for (uint32_t i = 0; i < presentModeCount; ++i)
    {
        //printf("Surface present mode: %d\n", presentModes[i]);
        if (presentModes[i] == VK_PRESENT_MODE_MAILBOX_KHR) presentModeIndex = i;
    }

    // Create swap chain
    _surfaceFormat = formats[formatIndex];
    _presentMode = presentModes[presentModeIndex];

    VkSwapchainCreateInfoKHR createInfo
    {
        VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR,
        nullptr, 0, _surface,
        caps.maxImageCount > 0 ? (caps.minImageCount + caps.maxImageCount) / 2
                               : (caps.minImageCount + 1),
        _surfaceFormat.format,
        _surfaceFormat.colorSpace,
        caps.maxImageExtent,
        1,  // imageArrayLayers
        VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT,  // imageUsage
        VK_SHARING_MODE_EXCLUSIVE,  // imageSharingMode
        0, nullptr,  // queueFamilyIndex
        caps.currentTransform,
        VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR,  // compositeAlpha
        _presentMode,
        VK_TRUE, nullptr
    };

    result = vkCreateSwapchainKHR(_logicDevice, &createInfo, nullptr, &_swapChain);
    if (result != VK_SUCCESS)
    {
        printf("Failed to create swap chain.\n");
        return false;
    }

    uint32_t imageCount = 0;
    vkGetSwapchainImagesKHR(_logicDevice, _swapChain, &imageCount, nullptr);
    printf("Get swapchain image count: %d\n", imageCount);

    _swapChainImages.resize(imageCount);
    vkGetSwapchainImagesKHR(_logicDevice, _swapChain, &imageCount, _swapChainImages.data());
    return true;
}

bool VulkanApplication::createImageViews()
{
    VkImageViewCreateInfo createInfo
    {
        VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO,
        nullptr, 0,
        nullptr, VK_IMAGE_VIEW_TYPE_2D,
        _surfaceFormat.format,
        { VK_COMPONENT_SWIZZLE_IDENTITY, VK_COMPONENT_SWIZZLE_IDENTITY,
          VK_COMPONENT_SWIZZLE_IDENTITY, VK_COMPONENT_SWIZZLE_IDENTITY },
        { VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1 }
    };

    _swapchainImageViews.resize(_swapChainImages.size());
    for (size_t i = 0; i < _swapchainImageViews.size(); ++i)
    {
        createInfo.image = _swapChainImages[i];
        VkResult result = vkCreateImageView(
            _logicDevice, &createInfo, nullptr, &_swapchainImageViews[i]);
        if (result != VK_SUCCESS)
        {
            printf("Failed to create image view: %zd.\n", i);
            return false;
        }
    }
    return true;
}

bool VulkanApplication::createCommandBuffer()
{
    // Create command pool
    if (_commandPool == nullptr)
    {
        VkCommandPoolCreateInfo poolInfo
        {
            VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO,
            nullptr,
            VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT,
            0  // queueFamilyIndex
        };

        VkResult result = vkCreateCommandPool(_logicDevice, &poolInfo, nullptr, &_commandPool);
        if (result != VK_SUCCESS)
        {
            printf("Failed to create command pool.\n");
            return false;
        }
    }

    // Create command buffer
    VkCommandBufferAllocateInfo allocInfo
    {
        VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO,
        nullptr, _commandPool,
        VK_COMMAND_BUFFER_LEVEL_PRIMARY, 1
    };

    VkResult result = vkAllocateCommandBuffers(_logicDevice, &allocInfo, &_commandBuffer);
    if (result != VK_SUCCESS)
    {
        printf("Failed to create command buffer.\n");
        return false;
    }
    return true;
}

bool VulkanApplication::createRenderPass()
{
    VkAttachmentDescription colorAttachment
    {
        0, _surfaceFormat.format,
        VK_SAMPLE_COUNT_1_BIT,
        VK_ATTACHMENT_LOAD_OP_CLEAR,       // Load/Store
        VK_ATTACHMENT_STORE_OP_STORE,
        VK_ATTACHMENT_LOAD_OP_DONT_CARE,   // stencilLoad/Store
        VK_ATTACHMENT_STORE_OP_DONT_CARE,
        VK_IMAGE_LAYOUT_UNDEFINED,         // initialLayout
        VK_IMAGE_LAYOUT_PRESENT_SRC_KHR    // finalLayout
    };

    VkAttachmentReference colorAttachmentRef
    {
        0, VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL
    };

    VkSubpassDescription subpass
    {
        0, VK_PIPELINE_BIND_POINT_GRAPHICS,
        0, nullptr,              // pInputAttachments
        1, &colorAttachmentRef,  // pColorAttachments
        nullptr,                 // pResolveAttachments
        nullptr,                 // pDepthStencilAttachment
        0, nullptr,              // pPreserveAttachments
    };

    VkRenderPassCreateInfo renderPassInfo
    {
        VK_STRUCTURE_TYPE_RENDER_PASS_CREATE_INFO,
        nullptr, 0,
        1, &colorAttachment,  // pAttachments
        1, &subpass,          // pSubpasses
        0, nullptr            // pDependencies
    };

    VkResult result = vkCreateRenderPass(_logicDevice, &renderPassInfo, nullptr, &_renderPass);
    if (result != VK_SUCCESS)
    {
        printf("Failed to create render-pass.\n");
        return false;
    }
    return true;
}

bool VulkanApplication::createFrameBuffers(int w, int h)
{
    _framebuffers.resize(_swapchainImageViews.size());
    for (size_t i = 0; i < _framebuffers.size(); ++i)
    {
        VkImageView attachments[] = { _swapchainImageViews[i] };
        VkFramebufferCreateInfo framebufferInfo
        {
            VK_STRUCTURE_TYPE_FRAMEBUFFER_CREATE_INFO,
            nullptr, 0,
            _renderPass,
            1, attachments,
            (uint32_t)w, (uint32_t)h, 1
        };
        vkCreateFramebuffer(_logicDevice, &framebufferInfo, nullptr, &_framebuffers[i]);
    }
    return true;
}

VkFence VulkanApplication::getOrCreateFence(const std::string& name)
{
    if (_fences.find(name) == _fences.end())
    {
        VkFenceCreateInfo fenceInfo
        {
            VK_STRUCTURE_TYPE_FENCE_CREATE_INFO,
            nullptr, VK_FENCE_CREATE_SIGNALED_BIT
        };
        vkCreateFence(_logicDevice, &fenceInfo, nullptr, &_fences[name]);
    }
    return _fences[name];
}

VkSemaphore VulkanApplication::getOrCreateSemaphore(const std::string& name)
{
    if (_semaphores.find(name) == _semaphores.end())
    {
        VkSemaphoreCreateInfo semaphoreInfo
        {
            VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO,
            nullptr, 0
        };
        vkCreateSemaphore(_logicDevice, &semaphoreInfo, nullptr, &_semaphores[name]);
    }
    return _semaphores[name];
}

void VulkanApplication::submitAndPresent(VkSemaphore waitImage, VkSemaphore waitSubmission,
                                         VkFence fenceSubmission, uint32_t imageIndex)
{
    VkPipelineStageFlags waitStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT;
    VkSwapchainKHR swapChains[] = { _swapChain };

    VkQueue queue = nullptr;
    vkGetDeviceQueue(_logicDevice, 0, 0, &queue);

    VkSubmitInfo submitInfo
    {
        VK_STRUCTURE_TYPE_SUBMIT_INFO,
        nullptr,
        1, &waitImage,       // pWaitSemaphores
        &waitStageMask,      // pWaitDstStageMask
        1, &_commandBuffer,
        1, &waitSubmission   // pSignalSemaphores
    };
    vkQueueSubmit(queue, 1, &submitInfo, fenceSubmission);

    VkPresentInfoKHR presentInfo
    {
        VK_STRUCTURE_TYPE_PRESENT_INFO_KHR,
        nullptr,
        1, &waitSubmission,   // pWaitSemaphores
        1, swapChains,        // pSwapchains
        &imageIndex,          // pImageIndices
        nullptr
    };
    vkQueuePresentKHR(queue, &presentInfo);
}